#include <iostream>
#include <string.h>

#include "packet_capture.h"

// struct pcap_pkthdr {
//     struct timeval ts; /* time stamp */
//     bpf_u_int32 caplen; /* length of portion present */
//     bpf_u_int32 len; /* length of this packet (off wire) */
// };

packet_capture::packet_capture(const std::string &device,
                               const std::string &filename)
    : handle_(nullptr)
    , dumper_(nullptr)
    , device_(device)
    , filename_(filename)
    , capturing_(false) {
    memset(errbuf_, 0, sizeof(errbuf_));
}

packet_capture::~packet_capture() { clear(); }

void packet_capture::start() {
    if (capturing_.load()) {  // 正在运行则返回
        return;
    }

    thread_ = std::thread(&packet_capture::capture_loop, this);

    {  // 等待捕获线程启动并准备好工作
        std::unique_lock<std::mutex> lk(cv_m_);
        start_cv_.wait(lk, [this] { return capturing_.load(); });
    }
}

void packet_capture::stop() {
    if (capturing_.load()) {
        pcap_breakloop(handle_);  // 使用信号通知停止抓包循环
        std::unique_lock<std::mutex> lk(cv_m_);
        // cv_等lambda表达式返回true
        cv_.wait(lk, [this] { return !capturing_.load(); });
    }
    if (thread_.joinable()) {
        thread_.join();
    }
}

void packet_capture::packet_handler(u_char *userData,
                                    const struct pcap_pkthdr *pkthdr,
                                    const u_char *packet) {
    // 写入数据包到文件中
    pcap_dump(userData, pkthdr, packet);
}

void packet_capture::capture_loop() {
    int snapshot_len = 65535;  // 抓包的最大包长度
    int promiscuous = 1;    // 1 启用promiscuous(混杂模式)抓包, 0禁用
    int timeout_ms = 1000;  // 读取网卡数据包的超时时间,单位毫秒

    // 打开网络监听会话句柄
    handle_ = pcap_open_live(device_.c_str(), snapshot_len, promiscuous,
                             timeout_ms, errbuf_);
    if (handle_ == nullptr) {
        std::cerr << "pcap_open_live() failed: " << errbuf_ << std::endl;
        return;
    }

    // 打开存储文件
    dumper_ = pcap_dump_open(handle_, filename_.c_str());
    if (dumper_ == nullptr) {
        std::cerr << "pcap_dump_open() failed: " << pcap_geterr(handle_)
                  << std::endl;
        return;
    }

    {
        std::unique_lock<std::mutex> lk(cv_m_);
        capturing_.store(true);
        start_cv_.notify_one();  // 通知主线程捕获线程已经准备好
    }
    // 开启抓包循环, 写入到dumper_中
    printf("Start capturing packets.\n");
    pcap_loop(handle_, 0, packet_handler, reinterpret_cast<u_char *>(dumper_));
    {
        std::unique_lock<std::mutex> lk(cv_m_);
        capturing_.store(false);
        cv_.notify_all();
    }

    // 清理抓包句柄
    clear();
    printf("Capture stopped, and file is saved as %s\n", filename_.c_str());
}

void packet_capture::clear() {
    if (handle_) {
        pcap_close(handle_);
        handle_ = nullptr;
    }
    if (dumper_) {
        pcap_dump_close(dumper_);
        dumper_ = nullptr;
    }
}

///////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////
packet_capture_guard::packet_capture_guard(const std::string &device,
                                           const std::string &filename) {
    obj_ = std::make_shared<packet_capture>(device, filename);
    obj_->start();
}

packet_capture_guard::~packet_capture_guard() { obj_->stop(); }
